from EXP3_Kmeans.gendata import show2dscatter,rawdata,show3dscatter,showNtime,showRet,show3dNtime,show3dRet
import  numpy as np
import random

data = rawdata()
show3dscatter(data)
k = 8
retdata = []
ret2data = []
for _ in range(k):
    retdata.append([])
    ret2data.append([])




gap = 1
last = 0
cur = 1
centers = []
newcenters = []

_i = k
rndset = []
while _i:

    idx = random.randint(0,len(data)-1)
    if idx in rndset:
        continue
    _i-=1
    rndset.append(idx)
    centers.append(data[idx])
    print("初始中心",data[idx])
    newcenters.append(np.array([0,0,0]))
centers = np.array(centers)
newcenters = np.array(newcenters)
n = 0
while True:

    sum = np.zeros((k,3))
    cnt = np.zeros(k)
    for i in range(len(data)):
        dis = [ np.linalg.norm(centers[_] - data[i], 2) for _ in range(k)]
        #print(data[i],dis)
        idx_min = np.argmin(dis)
        sum[idx_min] += data[i]
        retdata[idx_min].append(data[i])
        cnt[idx_min] +=1
    for _ in range(k):
        assert cnt[_] != 0,"某个类计数为0"
        newcenters[_] = sum[_]/cnt[_]
    if (newcenters == centers).all():
        print(n + 1, "-th iteration found ok centers")
        for _ in range(k):
            ret2data[_] = np.array(retdata[_])
        show3dRet(ret2data, k)
        break
    else:
        print(n+1,"-th iteration")
        for _ in range(k):
            ret2data[_] = np.array(retdata[_])
        show3dNtime(ret2data,k,n+1)
        n+=1
        centers = np.copy(newcenters)
        retdata.clear()
        for _ in range(k):
            retdata.append([])

